# coding: utf-8
######################################### tuhsare自动延迟下载 ####################################################
# tsettinguhsare自动延迟下载，防止频繁调取数据是报错
import time
import logging
import logging.handlers
import tushare as ts
import pandas as pd
import json
from dateutil.parser import parse
from DB.OthersDB import getTradeDays
# 请根据自己的情况填写ts的token

class TuShare:
    '''
    tushare服务接口自动重试封装类，能够在接口超时情况下自动等待1秒然后再次发起请求，
    无限次重试下去，直到返回结果或者达到最大重试次数。
    '''

    def __init__(self, token, logger=None, max_retry=0):
        """构造函数，token：tushare的token；logger：日志对象，可以不传；
        max_retry：最大重试次数，默认为0意为无限重试，建议用10以上100以内。"""
        self.token = token
        if not logger:
            logger = logging.getLogger('TuShare')
            # CRITICAL > ERROR > WARNING > INFO > DEBUG > NOTSET
            logger.setLevel(logging.DEBUG)
            formatter = logging.Formatter(
                '%(asctime)s %(name)s %(pathname)s:%(lineno)d %(funcName)s %(levelname)s %(message)s'
            )
            handler = logging.StreamHandler()
            handler.setFormatter(formatter)
            handler.setLevel(logging.DEBUG)
            logger.addHandler(handler)
        self.logger = logger
        self.max_retry = max_retry
        ts.set_token(token)
        self.pro = ts.pro_api()


    def __getattr__(self, name):
        def wrapper(*args, **kwargs):
            i = 0
            while True:
                try:
                    if name == 'pro_bar':
                        m = getattr(ts, name, None)
                    else:
                        m = getattr(self.pro, name, None)
                    if m is None:
                        self.logger.error('Attribute %s does not exist.', name)
                        return None
                    else:
                        return m(*args, **kwargs)
                except (Exception):
                    if self.max_retry > 0 and i >= self.max_retry:
                        raise
                    self.logger.exception(
                        'TuShare exec %s failed, args:%s, kwargs:%s, try again.',
                        name, args, kwargs)
                    time.sleep(1)
                i += 1

        return wrapper


######################################### 公用 ####################################################
# 绕过查询限制
def distributed_query(query_func_name,
                      symbol,
                      start_date,
                      end_date,
                      fields,
                      limit=3000):

    n_symbols = len(symbol.split(','))
    dates = query_trade_dates(start_date, end_date)
    n_days = len(dates)

    if n_symbols * n_days > limit:
        n = limit // n_symbols

        df_list = []
        i = 0
        pos1, pos2 = n * i, n * (i + 1) - 1

        while pos2 < n_days:

            df = query_func_name(
                ts_code=symbol,
                start_date=dates[pos1],
                end_date=dates[pos2],
                fields=fields)
            df_list.append(df)
            i += 1
            pos1, pos2 = n * i, n * (i + 1) - 1
        if pos1 < n_days:
            df = query_func_name(
                ts_code=symbol,
                start_date=dates[pos1],
                end_date=dates[-1],
                fields=fields)
            df_list.append(df)
        df = pd.concat(df_list, axis=0)
    else:
        df = query_func_name(
            ts_code=symbol,
            start_date=start_date,
            end_date=end_date,
            fields=fields)
    return df


# ts的日历需要处理一下才会返回成交日列表
def query_trade_dates(my_ts,start_date: str, end_date: str) -> list:
    #   use JQ trade day func
    start_date = parse(start_date).strftime('%Y-%m-%d')
    end_date = parse(end_date).strftime('%Y-%m-%d')

    trd_l = my_ts.trade_cal(is_open='1')['cal_date'].values.to_list()

    def trans_trd(d):
        d.strftime('%Y-%m-%d')

    trd_l = [trans_trd(i) for i in trd_l]
    return trd_l


# 设置 Tushare Pro 的 token 并初始化接口

def get_trade_dates(pro, end_date, count):
    # 获取交易日历
    trade_cal = pro.trade_cal(exchange='', start_date='20000101', end_date=end_date)

    # 筛选出交易日
    trade_dates = trade_cal[trade_cal['is_open'] == 1]['cal_date']

    # 将交易日转换为 datetime 类型
    trade_dates = pd.to_datetime(trade_dates, format='%Y%m%d')

    # 获取截止日期之前的 count 个交易日
    end_date_dt = pd.to_datetime(end_date, format='%Y%m%d')
    recent_trade_dates = trade_dates[trade_dates <= end_date_dt].head(count)

    # 将 datetime 类型转换回字符串
    recent_trade_dates = recent_trade_dates.dt.strftime('%Y%m%d').tolist()

    return recent_trade_dates


# if __name__ == '__main__':
#     # 示例：获取截至 20231031 的 10 个交易日
#     end_date = '20231031'
#     count = 10
#
#     tushare_token = '4b17847b701af3b6cd85fd6b7fc43c3805c37f1ed8fd18353b46f140'
#     my_ts = TuShare(tushare_token, max_retry=10)
#
#     recent_trade_dates = get_trade_dates(my_ts,end_date, count)
#     print(recent_trade_dates)






